#include "waferfrontend.hpp"

#include "clang/AST/AST.h"
#include "clang/AST/ASTConsumer.h"
#include "clang/AST/RecursiveASTVisitor.h"
#include "clang/Frontend/ASTConsumers.h"
#include "clang/Frontend/CompilerInstance.h"
#include "clang/Frontend/FrontendActions.h"
#include "clang/Rewrite/Core/Rewriter.h"
#include "clang/Tooling/CommonOptionsParser.h"
#include "clang/Tooling/Tooling.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/TargetParser/Host.h"

#include <fstream>

#include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVMPass.h"
#include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
#include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h"
#include "mlir/ExecutionEngine/ExecutionEngine.h"
#include "mlir/Parser/Parser.h"
#include "mlir/Target/LLVMIR/Dialect/All.h"
#include "mlir/Target/LLVMIR/Export.h"
#include "llvm/IR/Module.h"
#include "llvm/Support/TargetSelect.h"

#include "core/JsonHardwareInfo.hpp"

#include "graphSingleton.hpp"
#include "graph_node.pb.h"
#include <google/protobuf/util/json_util.h>

using namespace llvm;
using namespace llvm::opt;
using namespace mlir;
using namespace clang::driver;
using namespace clang::tooling;

// Begin command option

extern int cc1_main(ArrayRef<const char *> Argv, const char *Argv0,
                    void *MainAddr);
extern int cc1as_main(ArrayRef<const char *> Argv, const char *Argv0,
                      void *MainAddr);
extern int cc1gen_reproducer_main(ArrayRef<const char *> Argv,
                                  const char *Argv0, void *MainAddr);

static cl::OptionCategory toolOptions("clang to mlir - tool options");
static cl::list<std::string> includeDirs("I", cl::desc("include search path"),
                                         cl::cat(toolOptions));

static cl::opt<bool> PrintDebugInfo("print-debug-info", cl::init(false),
                                    cl::desc("Print debug info from MLIR"));

static cl::opt<bool> memRefABI("memref-abi", cl::init(true),
                               cl::desc("Use memrefs when possible"));

static cl::opt<bool> LowerCryptoToX86("lower-crypto-to-x86", cl::init(false),
                                      cl::desc("Use lower-crypto-to-x86 pass"));

static cl::opt<bool> LowerCryptoToX86WithoutSHA1ISA(
    "lower-cryptosha1-without-SHA1ISA", cl::init(false),
    cl::desc("Use lower-cryptosha1-without-SHA1ISA pass"));

static cl::opt<bool>
    LowerCryptoMultiToX86("lower-crypto-multi-x86", cl::init(false),
                          cl::desc("Use lower-crypto-multi-x86 pass"));

static cl::opt<bool> LowerCryptoToGpu("lower-crypto-to-gpu", cl::init(false),
                                      cl::desc("Use lower-crypto-to-gpu pass"));

static cl::opt<bool> ConvertSCFToCF("convert-scf-to-cf", cl::init(false),
                                    cl::desc("Convert SCF to CF"));

static cl::opt<bool> ConvertVectorToLLVM("convert-vector-to-llvm",
                                         cl::init(false),
                                         cl::desc("Convert Vector to LLVM"));

static cl::opt<bool> ConvertMemRefToLLVM("finalize-memref-to-llvm",
                                         cl::init(false),
                                         cl::desc("Convert MemRef to LLVM"));

static cl::opt<bool> ConvertArithToLLVM("convert-arith-to-llvm",
                                        cl::init(false),
                                        cl::desc("Convert Arith to LLVM"));

static cl::opt<bool>
    ReconcileUnrealizedCasts("reconcile-unrealized-casts", cl::init(false),
                             cl::desc("Reconcile Unrealized Casts"));

static cl::opt<bool> ConvertFuncToLLVM("convert-func-to-llvm", cl::init(false),
                                       cl::desc("Convert Func to LLVM"));

static cl::opt<bool> WaferToLLVMIR("wafer-to-llvmir", cl::init(false),
                                   cl::desc("Convert MLIR to LLVM IR"));

static cl::opt<bool> PTXTarget("ptx", cl::init(false),
                               cl::desc("generate gpu ptx"));

static cl::opt<bool> EmitImmMLIR("emit-mlir", cl::init(false),
                                 cl::desc("Emit mlir"));

cl::opt<bool> CStyleMemRef("c-style-memref", cl::init(true),
                           cl::desc("Use c style memrefs when possible"));

static cl::opt<bool> SSTTarget("sst", cl::init(false),
                               cl::desc("Convert Crypto to SST"));

static cl::opt<std::string> Output("o", cl::desc("output file"),
                                   cl::value_desc("filename"));

static cl::opt<std::string> JsonFilePath("JsonFilePath",
                                         cl::desc("JsonFilePath"),
                                         cl::value_desc("JsonFilePath"));

static cl::opt<std::string> GraphOutputPath("GraphOutputPath",
                                            cl::desc("GraphOutputPath"),
                                            cl::value_desc("GraphOutputPath"));

static cl::opt<bool> LowerCryptoToCGRA("lower-crypto-to-cgra", cl::init(false), cl::desc("Use lower-crypto-to-cgra pass"));

static cl::opt<bool> LowerCryptoToFPGA("lower-crypto-to-fpga", cl::init(false), cl::desc("Use lower-crypto-to-fpga pass"));
// End command option 

void WaferConsumer::HandleTranslationUnit(clang::ASTContext &Context) {
  WaferVisitor Visitor(&Context, *this, module);
  Visitor.TraverseDecl(Context.getTranslationUnitDecl());
}

bool WaferVisitor::VisitCallExpr(clang::CallExpr *expr) {
  const clang::FunctionDecl *funcDecl = expr->getDirectCallee();
  if (funcDecl && wafer_guard::WaferParallelStat) {
    std::string funcName = funcDecl->getNameAsString();
    // addNode
    if (funcName == "SHA1_MULTI_OP" || funcName == "SHA1_OP")
      GraphSingleton::addNode(Crypto::SHA1);
    else if (funcName == "MD5_MULTI_OP" || funcName == "MD5_OP")
      GraphSingleton::addNode(Crypto::MD5);
    else if (funcName == "SM3_MULTI_OP" || funcName == "SM3_OP")
      GraphSingleton::addNode(Crypto::SM3);
    else if (funcName == "AES_MULTI_OP" || funcName == "AES_OP")
      GraphSingleton::addNode(Crypto::AES);

    if (funcName == "SHA1_MULTI_OP") {
      const clang::SourceManager &SM = Context->getSourceManager();
      const clang::SourceLocation LocStart = expr->getBeginLoc();
      unsigned LineStart = SM.getSpellingLineNumber(LocStart);

      if (isNear(LineStart)) {
        const char *mlirString = R"(module {
          func.func @CUDA(%arg0: !llvm.ptr, %arg1: i32, %arg2: !llvm.ptr, %arg3: i32) {
            %0 = llvm.load %arg0 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
            %1 = llvm.load %arg2 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
            %2 = builtin.unrealized_conversion_cast %0 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
            %3 = builtin.unrealized_conversion_cast %1 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
            "crypto.sha1multi"(%2, %arg1, %3, %arg3) : (memref<?xi8>, i32, memref<?xi8>, i32) -> ()
            return
          }
        }
        )";
        mlir::OwningOpRef<mlir::ModuleOp> newModule =
            mlir::parseSourceString<mlir::ModuleOp>(mlirString,
                                                    builder.getContext());
        for (auto &op : newModule.get().getOps()) {
          module.get().push_back(builder.insert(op.clone()));
        }
      }

    } else if (funcName == "SHA1_OP") {
      const clang::SourceManager &SM = Context->getSourceManager();
      const clang::SourceLocation LocStart = expr->getBeginLoc();
      unsigned LineStart = SM.getSpellingLineNumber(LocStart);

      if (isNear(LineStart)) {
        const char *mlirString = R"(module {
          func.func @SHA1_OP(%arg0: !llvm.ptr, %arg1: i32, %arg2: !llvm.ptr) {
            %0 = llvm.load %arg0 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
            %1 = llvm.load %arg2 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
            %2 = builtin.unrealized_conversion_cast %0 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
            %3 = builtin.unrealized_conversion_cast %1 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
            "crypto.sha1"(%2, %arg1, %3) : (memref<?xi8>, i32, memref<?xi8>) -> ()
            return
          }
        })";
        mlir::OwningOpRef<mlir::ModuleOp> newModule =
            mlir::parseSourceString<mlir::ModuleOp>(mlirString,
                                                    builder.getContext());
        for (auto &op : newModule.get().getOps()) {
          module.get().push_back(builder.insert(op.clone()));
        }
      }
    } else if (funcName == "MD5_OP") {
      const clang::SourceManager &SM = Context->getSourceManager();
      const clang::SourceLocation LocStart = expr->getBeginLoc();
      unsigned LineStart = SM.getSpellingLineNumber(LocStart);

      if (isNear(LineStart)) {
        const char *mlirString;
        if (LowerCryptoToCGRA) {
          mlirString = R"(module {
            func.func @MD5_OP(%arg0: !llvm.ptr, %arg1: i32, %arg2: !llvm.ptr) {
            %0 = llvm.load %arg0 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
            %1 = llvm.load %arg2 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
            %2 = builtin.unrealized_conversion_cast %0 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
            %3 = builtin.unrealized_conversion_cast %1 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
            "crypto.md5"(%2, %arg1, %3) : (memref<?xi8>, i32, memref<?xi8>) -> ()
            return
            }
        })";
        } else {
          mlirString = R"(module {
            func.func @CUDA(%arg0: !llvm.ptr, %arg1: i32, %arg2: !llvm.ptr, %arg3: i32) {
              %0 = llvm.load %arg0 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
              %1 = llvm.load %arg2 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
              %2 = builtin.unrealized_conversion_cast %0 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
              %3 = builtin.unrealized_conversion_cast %1 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
              "crypto.md5multi"(%2, %arg1, %3, %arg3) : (memref<?xi8>, i32, memref<?xi8>, i32) -> ()
              return
            }
          }
          )";
        }
        mlir::OwningOpRef<mlir::ModuleOp> newModule =
            mlir::parseSourceString<mlir::ModuleOp>(mlirString,
                                                    builder.getContext());
        for (auto &op : newModule.get().getOps()) {
          module.get().push_back(builder.insert(op.clone()));
        }
      }
    } else if (funcName == "SM3_OP") {
      const clang::SourceManager &SM = Context->getSourceManager();
      const clang::SourceLocation LocStart = expr->getBeginLoc();
      unsigned LineStart = SM.getSpellingLineNumber(LocStart);

      if (isNear(LineStart)) {
        const char *mlirString;
        if (LowerCryptoToCGRA) {
          mlirString = R"(module {
            func.func @SM3_OP(%arg0: !llvm.ptr, %arg1: i32, %arg2: !llvm.ptr) {
            %0 = llvm.load %arg0 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
            %1 = llvm.load %arg2 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
            %2 = builtin.unrealized_conversion_cast %0 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
            %3 = builtin.unrealized_conversion_cast %1 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
            "crypto.sm3"(%2, %arg1, %3) : (memref<?xi8>, i32, memref<?xi8>) -> ()
            return
            }
          })";
        } else {
          mlirString = R"(module {
            func.func @CUDA(%arg0: !llvm.ptr, %arg1: i32, %arg2: !llvm.ptr, %arg3: i32) {
              %0 = llvm.load %arg0 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
              %1 = llvm.load %arg2 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
              %2 = builtin.unrealized_conversion_cast %0 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
              %3 = builtin.unrealized_conversion_cast %1 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
              "crypto.sm3multi"(%2, %arg1, %3, %arg3) : (memref<?xi8>, i32, memref<?xi8>, i32) -> ()
              return
            }
          }
          )";
        }
        mlir::OwningOpRef<mlir::ModuleOp> newModule =
            mlir::parseSourceString<mlir::ModuleOp>(mlirString,
                                                    builder.getContext());
        for (auto &op : newModule.get().getOps()) {
          module.get().push_back(builder.insert(op.clone()));
        }
      }
    } else if (funcName == "AES_OP") {
      const clang::SourceManager &SM = Context->getSourceManager();
      const clang::SourceLocation LocStart = expr->getBeginLoc();
      unsigned LineStart = SM.getSpellingLineNumber(LocStart);

      if (isNear(LineStart)) {
        const char *mlirString = R"(module {
          func.func @CUDA(%arg0: !llvm.ptr, %arg1: i32, %arg2: !llvm.ptr, %arg3: i32, %arg4: !llvm.ptr, %arg5: i32, %arg6 : i32) {
            %0 = llvm.load %arg0 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
            %1 = llvm.load %arg2 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
            %2 = llvm.load %arg4 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
            %3 = builtin.unrealized_conversion_cast %0 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
            %4 = builtin.unrealized_conversion_cast %1 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
            %5 = builtin.unrealized_conversion_cast %2 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
            "crypto.aesmulti"(%3, %arg1, %4, %arg3, %5, %arg5, %arg6) : (memref<?xi8>, i32, memref<?xi8>, i32, memref<?xi8>, i32, i32) -> ()
            return
          }
        }
        )";
        mlir::OwningOpRef<mlir::ModuleOp> newModule =
            mlir::parseSourceString<mlir::ModuleOp>(mlirString,
                                                    builder.getContext());
        for (auto &op : newModule.get().getOps()) {
          module.get().push_back(builder.insert(op.clone()));
        }
      }
    }

    if (funcName == "SHA1_CRACK_OP") {
      const clang::SourceManager &SM = Context->getSourceManager();
      const clang::SourceLocation LocStart = expr->getBeginLoc();
      unsigned LineStart = SM.getSpellingLineNumber(LocStart);

      if (isNear(LineStart)) {
        const char *mlirString = R"(module {
          func.func @CUDA(%hash: !llvm.ptr, %candidate: !llvm.ptr, %num: i32, %len: i32, %out: !llvm.ptr) {
            %0 = llvm.load %hash : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
            %1 = llvm.load %candidate : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
            %2 = llvm.load %out : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
            %3 = builtin.unrealized_conversion_cast %0 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
            %4 = builtin.unrealized_conversion_cast %1 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
            %5 = builtin.unrealized_conversion_cast %2 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
            %res = "crypto.sha1.crack" (%3, %4, %num, %len) : (memref<?xi8>, memref<?xi8>, i32, i32) -> (memref<?xi8>)
            memref.copy %res, %5 : memref<?xi8> to memref<?xi8>
            return
          }
        })";
        mlir::OwningOpRef<mlir::ModuleOp> newModule =
            mlir::parseSourceString<mlir::ModuleOp>(mlirString,
                                                    builder.getContext());
        for (auto &op : newModule.get().getOps()) {
          module.get().push_back(builder.insert(op.clone()));
        }
      }
    } else if (funcName == "MD5_CRACK_OP") {
      const clang::SourceManager &SM = Context->getSourceManager();
      const clang::SourceLocation LocStart = expr->getBeginLoc();
      unsigned LineStart = SM.getSpellingLineNumber(LocStart);

      if (isNear(LineStart)) {
        const char *mlirString = R"(module {
          func.func @CUDA(%hash: !llvm.ptr, %candidate: !llvm.ptr, %num: i32, %len: i32, %out: !llvm.ptr) {
            %0 = llvm.load %hash : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
            %1 = llvm.load %candidate : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
            %2 = llvm.load %out : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
            %3 = builtin.unrealized_conversion_cast %0 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
            %4 = builtin.unrealized_conversion_cast %1 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
            %5 = builtin.unrealized_conversion_cast %2 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
            %res = "crypto.md5.crack" (%3, %4, %num, %len) : (memref<?xi8>, memref<?xi8>, i32, i32) -> (memref<?xi8>)
            memref.copy %res, %5 : memref<?xi8> to memref<?xi8>
            return
          }
        })";
        mlir::OwningOpRef<mlir::ModuleOp> newModule =
            mlir::parseSourceString<mlir::ModuleOp>(mlirString,
                                                    builder.getContext());
        for (auto &op : newModule.get().getOps()) {
          module.get().push_back(builder.insert(op.clone()));
        }
      }
    } else if (funcName == "SM3_CRACK_OP") {
      const clang::SourceManager &SM = Context->getSourceManager();
      const clang::SourceLocation LocStart = expr->getBeginLoc();
      unsigned LineStart = SM.getSpellingLineNumber(LocStart);

      if (isNear(LineStart)) {
        const char *mlirString = R"(module {
          func.func @CUDA(%hash: !llvm.ptr, %candidate: !llvm.ptr, %num: i32, %len: i32, %out: !llvm.ptr) {
            %0 = llvm.load %hash : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
            %1 = llvm.load %candidate : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
            %2 = llvm.load %out : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
            %3 = builtin.unrealized_conversion_cast %0 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
            %4 = builtin.unrealized_conversion_cast %1 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
            %5 = builtin.unrealized_conversion_cast %2 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
            %res = "crypto.sm3.crack" (%3, %4, %num, %len) : (memref<?xi8>, memref<?xi8>, i32, i32) -> (memref<?xi8>)
            memref.copy %res, %5 : memref<?xi8> to memref<?xi8>
            return
          }
        })";
        mlir::OwningOpRef<mlir::ModuleOp> newModule =
            mlir::parseSourceString<mlir::ModuleOp>(mlirString,
                                                    builder.getContext());
        for (auto &op : newModule.get().getOps()) {
          module.get().push_back(builder.insert(op.clone()));
        }
      }
    } else if (funcName == "AES_CRACK_OP") {
      const clang::SourceManager &SM = Context->getSourceManager();
      const clang::SourceLocation LocStart = expr->getBeginLoc();
      unsigned LineStart = SM.getSpellingLineNumber(LocStart);

      if (isNear(LineStart)) {
        const char *mlirString = R"(module {
          func.func @CUDA(%arg0: !llvm.ptr, %arg1: i32, %arg2: !llvm.ptr, %arg3: i32, %arg4: !llvm.ptr, %arg5: i32, %arg6 : i32, %out: !llvm.ptr) {
            %0 = llvm.load %arg0 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
            %1 = llvm.load %arg2 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
            %2 = llvm.load %arg4 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>

            %_out = llvm.load %out : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
            
            %3 = builtin.unrealized_conversion_cast %0 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
            %4 = builtin.unrealized_conversion_cast %1 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>
            %5 = builtin.unrealized_conversion_cast %2 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>

            %6 = builtin.unrealized_conversion_cast %_out : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xi8>

            %res = "crypto.aes.crack"(%3, %arg1, %4, %arg3, %5, %arg5, %arg6) : (memref<?xi8>, i32, memref<?xi8>, i32, memref<?xi8>, i32, i32) -> (memref<?xi8>)
            memref.copy %res, %6 : memref<?xi8> to memref<?xi8>
            return
          }
        })";
        mlir::OwningOpRef<mlir::ModuleOp> newModule =
            mlir::parseSourceString<mlir::ModuleOp>(mlirString,
                                                    builder.getContext());
        for (auto &op : newModule.get().getOps()) {
          module.get().push_back(builder.insert(op.clone()));
        }
      }
    }
  }
  return true;
}

static bool parseMLIR(const char *Argv0, std::vector<std::string> filenames,
                      std::vector<std::string> includeDirs,
                      mlir::OwningOpRef<mlir::ModuleOp> &module) {

  IntrusiveRefCntPtr<DiagnosticIDs> DiagID(new DiagnosticIDs());
  IntrusiveRefCntPtr<DiagnosticOptions> DiagOpts = new DiagnosticOptions();
  TextDiagnosticBuffer *DiagsBuffer = new TextDiagnosticBuffer;
  DiagnosticsEngine Diags(DiagID, &*DiagOpts, DiagsBuffer);

  bool Success;
  const char *binary = Argv0;
  const std::unique_ptr<Driver> driver(
      new Driver(binary, sys::getDefaultTargetTriple(), Diags));
  mlirclang::ArgumentList Argv;
  Argv.push_back(binary);
  for (const auto &filename : filenames) {
    Argv.emplace_back(filename);
  }

  const std::unique_ptr<Compilation> compilation(
      driver->BuildCompilation(Argv.getArguments()));
  JobList &Jobs = compilation->getJobs();
  if (Jobs.size() < 1)
    return false;

  WaferAction Act(module);

  for (auto &job : Jobs) {
    std::unique_ptr<CompilerInstance> Clang(new CompilerInstance());

    Command *cmd = cast<Command>(&job);
    if (strcmp(cmd->getCreator().getName(), "clang"))
      return false;

    const ArgStringList *args = &cmd->getArguments();

    Success = CompilerInvocation::CreateFromArgs(Clang->getInvocation(), *args,
                                                 Diags);
    Clang->getInvocation().getFrontendOpts().DisableFree = false;

    void *GetExecutablePathVP = (void *)(intptr_t)GetExecutablePath;
    if (Clang->getHeaderSearchOpts().UseBuiltinIncludes &&
        Clang->getHeaderSearchOpts().ResourceDir.size() == 0)
      Clang->getHeaderSearchOpts().ResourceDir =
          CompilerInvocation::GetResourcesPath(Argv0, GetExecutablePathVP);

    for (const auto &dir : includeDirs) {
      Clang->getHeaderSearchOpts().AddPath(dir, frontend::System, false, false);
    }

    Clang->getInvocation().getFrontendOpts().DisableFree = false;

    // Create the actual diagnostics engine.
    Clang->createDiagnostics();
    if (!Clang->hasDiagnostics())
      return false;

    DiagsBuffer->FlushDiagnostics(Clang->getDiagnostics());
    if (!Success)
      return false;

    // Create the target instance.
    Clang->setTarget(TargetInfo::CreateTargetInfo(
        Clang->getDiagnostics(), Clang->getInvocation().TargetOpts));
    if (!Clang->hasTarget())
      return false;

    Clang->getTarget().adjust(Clang->getDiagnostics(), Clang->getLangOpts());

    // Adjust target options based on codegen options.
    Clang->getTarget().adjustTargetOptions(Clang->getCodeGenOpts(),
                                           Clang->getTargetOpts());

    for (const auto &FIF : Clang->getFrontendOpts().Inputs) {
      if (Clang->hasSourceManager() && !Act.isModelParsingAction())
        Clang->getSourceManager().clearIDTables();
      if (Act.BeginSourceFile(*Clang, FIF)) {

        llvm::Error err = Act.Execute();
        if (err) {
          llvm::errs() << "saw error: " << err << "\n";
          return false;
        }
        assert(Clang->hasSourceManager());

        Act.EndSourceFile();
      }
    }
  }
  return true;
}

std::unique_ptr<Pass> createLowerCryptoToHyperPass();
std::unique_ptr<Pass> createLowerPadMessagePass();
std::unique_ptr<Pass> createLowerCryptoToX86WithoutSHA1ISAPass();
std::unique_ptr<Pass> createSSTDeviceSchedulePass();
std::unique_ptr<Pass> createSSTLowerHyperPass();
std::unique_ptr<Pass> createLowerHyperSSTMemcpyPass();
std::unique_ptr<Pass> createSSTToLLVMPass();
std::unique_ptr<Pass> createLowerCryptoToGpuPass();
std::unique_ptr<Pass> createDeviceSchedulePass();
std::unique_ptr<Pass> createLowerCryptoToCGRAPass();
std::unique_ptr<Pass> createLowerCryptoToFPGAPass();

int main(int argc, char **argv) {
  GOOGLE_PROTOBUF_VERIFY_VERSION;
  SmallVector<const char *> LinkageArgs;
  SmallVector<const char *> MLIRArgs;
  std::string outputFilePath;

  {
    bool linkOnly = false;
    for (int i = 0; i < argc; i++) {
      StringRef ref(argv[i]);
      if (ref == "-Wl,--start-group")
        linkOnly = true;
      if (!linkOnly) {
        if (ref == "-fPIC" || ref == "-c" || ref.startswith("-fsanitize")) {
          LinkageArgs.push_back(argv[i]);
        } else if (ref == "-L" || ref == "-l") {
          LinkageArgs.push_back(argv[i]);
          i++;
          LinkageArgs.push_back(argv[i]);
        } else if (ref.startswith("-L") || ref.startswith("-l") ||
                   ref.startswith("-Wl")) {
          LinkageArgs.push_back(argv[i]);
        } else if (ref.startswith("-D")) {
          MLIRArgs.push_back("-D");
          MLIRArgs.push_back(&argv[i][2]);
        } else if (ref == "-g") {
          LinkageArgs.push_back(argv[i]);
        } else {
          MLIRArgs.push_back(argv[i]);
        }
      } else {
        LinkageArgs.push_back(argv[i]);
      }
      if (ref == "-Wl,--end-group")
        linkOnly = false;
    }
  }
  using namespace mlir;

  int size = MLIRArgs.size();
  const char **data = MLIRArgs.data();
  std::vector<std::string> files;
  {
    cl::list<std::string> inputFileName(cl::Positional, cl::OneOrMore,
                                        cl::desc("<Specify input file>"),
                                        cl::cat(toolOptions));
    cl::ParseCommandLineOptions(size, data);
    assert(inputFileName.size());
    for (auto inp : inputFileName) {
      std::ifstream inputFile(inp);
      if (!inputFile.good()) {
        outs() << "Not able to open file: " << inp << "\n";
        return -1;
      }
      files.push_back(inp);
    }
  }

  mlir::DialectRegistry registry;
  registerArmNeonDialectTranslation(registry);
  registerAMXDialectTranslation(registry);
  registerArmSVEDialectTranslation(registry);
  registerLLVMDialectTranslation(registry);
  registerNVVMDialectTranslation(registry);
  registerOpenACCDialectTranslation(registry);
  registerOpenMPDialectTranslation(registry);
  registerROCDLDialectTranslation(registry);
  registerX86VectorDialectTranslation(registry);
  registerGpuSerializeToCubinPass();
  registerAllToLLVMIRTranslations(registry);
  MLIRContext context(registry);

  context.disableMultithreading();
  context.loadDialect<func::FuncDialect>();
  context.loadDialect<memref::MemRefDialect>();
  context.loadDialect<vector::VectorDialect>();
  context.loadDialect<arith::ArithDialect>();
  context.loadDialect<scf::SCFDialect>();
  context.loadDialect<crypto::CryptoDialect>();
  context.loadDialect<gpu::GPUDialect>();
  context.loadDialect<hyper::HyperDialect>();
  context.loadDialect<sst::SSTDialect>();

  auto mlir_loc =
      mlir::ModuleOp::create(mlir::OpBuilder(&context).getUnknownLoc());
  mlir::OwningOpRef<mlir::ModuleOp> module(mlir_loc);
  mlir::registerLLVMDialectTranslation(*(module.get())->getContext());
  llvm::LLVMContext llvmContext;
  auto llvmModule = mlir::translateModuleToLLVMIR(module.get(), llvmContext);

  parseMLIR(argv[0], files, includeDirs, module);

  mlir::PassManager pm(&context);

  if (!JsonFilePath.empty()) {
    wafer::Root &HardwareInfo = wafer::Root::getInstance();
    HardwareInfo.initialize(JsonFilePath.getValue());
    // lower crypto to sst simulator
    if (SSTTarget) {
      // TODO: Automatically get ptx files
      pm.addPass(createLowerCryptoToHyperPass());
      pm.addPass(createSSTDeviceSchedulePass());
      pm.addPass(createLowerHyperSSTMemcpyPass());
      pm.addPass(createSSTLowerHyperPass());
      pm.addPass(createLowerPadMessagePass());
      pm.addPass(createLowerCryptoToX86WithoutSHA1ISAPass());
      pm.addPass(mlir::createConvertSCFToCFPass());
      pm.addPass(mlir::createArithToLLVMConversionPass());
      pm.addPass(mlir::createFinalizeMemRefToLLVMConversionPass());
      pm.addPass(mlir::createConvertVectorToLLVMPass());
      pm.addPass(createSSTToLLVMPass());
      pm.addPass(mlir::createConvertFuncToLLVMPass());
      pm.addPass(mlir::createReconcileUnrealizedCastsPass());

      if (mlir::failed(pm.run(module.get()))) {
        module->dump();
        return 1;
      }
    } else if (PTXTarget) {
      pm.addPass(createLowerCryptoToHyperPass());

      if (mlir::failed(pm.run(module.get()))) {
        module->dump();
        return 1;
      }
    } else {
      std::cerr << "error: we only support sst! please add option sst or "
                   "lower-crypto-to-ptx!\n";
    }
  }

  // LowerCryptoToX86 pass
  if (!LowerCryptoMultiToX86 && LowerCryptoToX86) {
    pm.addPass(mlir::createConvertSCFToCFPass());
    pm.addPass(mlir::createConvertVectorToLLVMPass());
    pm.addPass(mlir::createFinalizeMemRefToLLVMConversionPass());
    pm.addPass(mlir::createArithToLLVMConversionPass());
    pm.addPass(mlir::createConvertFuncToLLVMPass());
    pm.addPass(mlir::createReconcileUnrealizedCastsPass());
    if (mlir::failed(pm.run(module.get()))) {
      module->dump();
      return 1;
    }
  }
  // End LowerCryptoToX86 pass

  // Begin LowerCryptoToX86WithOutSHA1ISA pass
  if (!LowerCryptoMultiToX86 && LowerCryptoToX86WithoutSHA1ISA) {
    pm.addPass(createLowerCryptoToX86WithoutSHA1ISAPass());
    pm.addPass(mlir::createConvertSCFToCFPass());
    pm.addPass(mlir::createConvertVectorToLLVMPass());
    pm.addPass(mlir::createFinalizeMemRefToLLVMConversionPass());
    pm.addPass(mlir::createArithToLLVMConversionPass());
    pm.addPass(mlir::createConvertFuncToLLVMPass());
    pm.addPass(mlir::createReconcileUnrealizedCastsPass());
    if (mlir::failed(pm.run(module.get()))) {
      module->dump();
      return 1;
    }
  }
  // End LowerCryptoToX86WithOutSHA1ISA pass

  // LowerCryptoToGpu pass
  if (!LowerCryptoMultiToX86 && LowerCryptoToGpu) {
    // pm.addPass(createLowerCryptoToGpuPass());
    pm.addPass(mlir::createGpuKernelOutliningPass());
    pm.addPass(mlir::createConvertSCFToCFPass());
    pm.addPass(mlir::createArithToLLVMConversionPass());
    pm.addPass(mlir::createFinalizeMemRefToLLVMConversionPass());
    pm.addPass(mlir::createConvertVectorToLLVMPass());
    pm.addPass(mlir::createConvertFuncToLLVMPass());
    pm.addNestedPass<gpu::GPUModuleOp>(createStripDebugInfoPass());
    pm.addNestedPass<gpu::GPUModuleOp>(createLowerGpuOpsToNVVMOpsPass());
    pm.addNestedPass<gpu::GPUModuleOp>(createReconcileUnrealizedCastsPass());
    pm.addNestedPass<mlir::gpu::GPUModuleOp>(
        mlir::createGpuSerializeToCubinPass("nvptx64-nvidia-cuda", "sm_75",
                                            "+ptx60"));
    pm.addNestedPass<mlir::func::FuncOp>(mlir::createGpuAsyncRegionPass());
    pm.addPass(mlir::createGpuToLLVMConversionPass());
    pm.addPass(mlir::createReconcileUnrealizedCastsPass());
    if (mlir::failed(pm.run(module.get()))) {
      module->dump();
      return 1;
    }
  }
  // End LowerCryptoToGpu pass

  // LowerCryptoMultiToX86 pass
  if (LowerCryptoMultiToX86) {
    if (LowerCryptoToX86) {
      pm.addPass(mlir::createConvertSCFToCFPass());
      pm.addPass(mlir::createConvertVectorToLLVMPass());
      pm.addPass(mlir::createFinalizeMemRefToLLVMConversionPass());
      pm.addPass(mlir::createArithToLLVMConversionPass());
      pm.addPass(mlir::createConvertFuncToLLVMPass());
      pm.addPass(mlir::createReconcileUnrealizedCastsPass());
    }
    if (LowerCryptoToX86WithoutSHA1ISA) {
      pm.addPass(createLowerCryptoToX86WithoutSHA1ISAPass());
      pm.addPass(mlir::createConvertSCFToCFPass());
      pm.addPass(mlir::createConvertVectorToLLVMPass());
      pm.addPass(mlir::createFinalizeMemRefToLLVMConversionPass());
      pm.addPass(mlir::createArithToLLVMConversionPass());
      pm.addPass(mlir::createConvertFuncToLLVMPass());
      pm.addPass(mlir::createReconcileUnrealizedCastsPass());
    }
    if (mlir::failed(pm.run(module.get()))) {
      module->dump();
      return 1;
    }
  }
  // End LowerCryptoMultiToX86 pass

  // Begin LowerCryptoToCGRA pass
  if (LowerCryptoToCGRA) {
    pm.addPass(createLowerCryptoToCGRAPass());
    pm.addPass(mlir::createFinalizeMemRefToLLVMConversionPass());
    pm.addPass(mlir::createConvertFuncToLLVMPass());
    pm.addPass(mlir::createReconcileUnrealizedCastsPass());
    if (mlir::failed(pm.run(module.get()))) {
      module->dump();
      return 1;
    }
  }
// End LowerCryptoToCGRA pass 

// Begin LowerCryptoToFPGA pass
if(LowerCryptoToFPGA) {
  pm.addPass(createLowerCryptoToFPGAPass());
  pm.addPass(mlir::createFinalizeMemRefToLLVMConversionPass());
  pm.addPass(mlir::createConvertFuncToLLVMPass());
  pm.addPass(mlir::createReconcileUnrealizedCastsPass());
  if (mlir::failed(pm.run(module.get()))) {
    module->dump();
    return 1;
  }
}
// End LowerCryptoToFPGA pass

  OpPrintingFlags flags;
  if (PrintDebugInfo)
    flags.enableDebugInfo(/*pretty=*/false);

  // Set outputpath
  if (!Output.empty()) {
    std::error_code errorCode;
    std::string outputFilePath = Output;

    llvm::raw_fd_ostream outputFile(outputFilePath, errorCode,
                                    llvm::sys::fs::OF_Text);
    if (errorCode) {
      llvm::errs() << "Failed to open output file: " << errorCode.message()
                   << "\n";
      return 1;
    }

    if (WaferToLLVMIR) {
      mlir::registerLLVMDialectTranslation(*module->getContext());
      llvm::LLVMContext llvmContext;
      auto llvmModule =
          mlir::translateModuleToLLVMIR(module.get(), llvmContext);
      if (!llvmModule) {
        llvm::errs() << "Failed to emit LLVM IR\n";
        return -1;
      }

      // Initialize LLVM targets.
      llvm::InitializeNativeTarget();
      llvm::InitializeNativeTargetAsmPrinter();

      auto tmBuilderOrError = llvm::orc::JITTargetMachineBuilder::detectHost();
      if (!tmBuilderOrError)
        return -1;

      auto tmOrError = tmBuilderOrError->createTargetMachine();
      if (!tmOrError)
        return -1;

      mlir::ExecutionEngine::setupTargetTripleAndDataLayout(
          llvmModule.get(), tmOrError.get().get());

      llvmModule->print(outputFile, nullptr);
    } else
      module->print(outputFile, flags);
  }

  // Emit immediate mlir code
  if (EmitImmMLIR)
    module->print(llvm::errs(), flags);

  if (!GraphOutputPath.empty()) {
    GraphSingleton::dump(GraphOutputPath.getValue());
  }
  return 0;
}
